package org.springframework.security.web.authentication.rememberme;

import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.security.core.userdetails.UserDetailsService;
import org.springframework.security.crypto.codec.Hex;
import org.springframework.security.crypto.codec.Utf8;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;

import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;
import java.util.Date;

/**
 * @author Dillon
 * @date 2024/7/12
 * @slogan 致敬大师 致敬未来的你
 * @desc 基于token的方式实现的记住我逻辑
 */
public class TokenBasedRememberMeServices extends AbstractRememberMeServices {

	/**
	 * 记住我 默认匹配算法
	 */
	private static final RememberMeTokenAlgorithm DEFAULT_MATCHING_ALGORITHM = RememberMeTokenAlgorithm.SHA256;

	/**
	 * 记住我默认加密算法
	 */
	private static final RememberMeTokenAlgorithm DEFAULT_ENCODING_ALGORITHM = RememberMeTokenAlgorithm.SHA256;

	/**
	 * 记住我加密算法
	 */
	private final RememberMeTokenAlgorithm encodingAlgorithm;

	/**
	 * 记住我匹配算法
	 */
	private RememberMeTokenAlgorithm matchingAlgorithm = DEFAULT_MATCHING_ALGORITHM;

	/**
	 * 构造函数
	 *
	 * @param key                记住我 key
	 * @param userDetailsService 数据源获取实现类
	 */
	public TokenBasedRememberMeServices(String key, UserDetailsService userDetailsService) {
		this(key, userDetailsService, DEFAULT_ENCODING_ALGORITHM);
	}

	/**
	 * 构造函数
	 *
	 * @param key                记住我 key
	 * @param userDetailsService 数据源获取实现类
	 * @param encodingAlgorithm  记住我加密算法
	 */
	public TokenBasedRememberMeServices(String key, UserDetailsService userDetailsService,
			RememberMeTokenAlgorithm encodingAlgorithm) {
		super(key, userDetailsService);
		Assert.notNull(encodingAlgorithm, "算法不能为空");
		this.encodingAlgorithm = encodingAlgorithm;
	}

	/**
	 * 验证token 并获取登录的user对象
	 *
	 * @param cookieTokens 解密后的cookie数组
	 * @param request      请求对象
	 * @param response     响应对象
	 * @return user对象
	 */
	@Override
	protected UserDetails processAutoLoginCookie(String[] cookieTokens, HttpServletRequest request, HttpServletResponse response) {
		// token长度验证
		if (!isValidCookieTokensLength(cookieTokens)) {
			throw new InvalidCookieException("cookie数组长度不对 当前数组为'" + Arrays.asList(cookieTokens) + "'");
		}
		// 获取并校验超时使劲啊
		long tokenExpiryTime = getTokenExpiryTime(cookieTokens);
		if (isTokenExpired(tokenExpiryTime)) {
			throw new InvalidCookieException("cookie超时时间为 '" + new Date(tokenExpiryTime) + "'; 当前时间为 '" + new Date() + "')");
		}
		// 获取登录用户并校验是否获取到
		UserDetails userDetails = getUserDetailsService().loadUserByUsername(cookieTokens[0]);
		Assert.notNull(userDetails, () -> "数据源 " + getUserDetailsService() + " 未查到用户 " + cookieTokens[0]);

		// 加签并比较签名
		String actualTokenSignature = cookieTokens[2];
		RememberMeTokenAlgorithm actualAlgorithm = this.matchingAlgorithm;
		// I如果包含签名算法，则这里需要取最后一个值 说明第三个为算法
		if (cookieTokens.length == 4) {
			actualTokenSignature = cookieTokens[3];
			actualAlgorithm = RememberMeTokenAlgorithm.valueOf(cookieTokens[2]);
		}
		// 加签
		String expectedTokenSignature = makeTokenSignature(tokenExpiryTime, userDetails.getUsername(),
				userDetails.getPassword(), actualAlgorithm);
		// 验签
		if (!equals(expectedTokenSignature, actualTokenSignature)) {
			throw new InvalidCookieException("cookie签名为 '" + actualTokenSignature + "' 但期望签名为 '" + expectedTokenSignature + "'");
		}
		return userDetails;
	}

	/**
	 * 验证cookie数组长度
	 * 3 名称 超时时间 签名
	 * 4 名称 超时时间 算法 签名
	 *
	 * @param cookieTokens cookie数组
	 * @return 是否通过
	 */
	private boolean isValidCookieTokensLength(String[] cookieTokens) {
		return cookieTokens.length == 3 || cookieTokens.length == 4;
	}

	/**
	 * 获取超时时间
	 *
	 * @param cookieTokens cookie数组
	 * @return
	 */
	private long getTokenExpiryTime(String[] cookieTokens) {
		try {
			return Long.parseLong(cookieTokens[1]);
		} catch (NumberFormatException nfe) {
			throw new InvalidCookieException("cookie 未包含一个正确的超时时间 '" + cookieTokens[1] + "')");
		}
	}

	/**
	 * 加签
	 *
	 * @param tokenExpiryTime 超时时间
	 * @param username        用户名
	 * @param password        密码
	 * @return 加签结果
	 */
	protected String makeTokenSignature(long tokenExpiryTime, String username, String password) {
		String data = username + ":" + tokenExpiryTime + ":" + password + ":" + getKey();
		try {
			MessageDigest digest = MessageDigest.getInstance(this.encodingAlgorithm.getDigestAlgorithm());
			return new String(Hex.encode(digest.digest(data.getBytes())));
		} catch (NoSuchAlgorithmException ex) {
			throw new IllegalStateException("No " + this.encodingAlgorithm.name() + " algorithm available!");
		}
	}

	/**
	 * 加签
	 *
	 * @param tokenExpiryTime 超时时间
	 * @param username        用户名
	 * @param password        密码
	 * @param algorithm       算法
	 * @return 加签结果
	 */
	protected String makeTokenSignature(long tokenExpiryTime, String username, String password,
			RememberMeTokenAlgorithm algorithm) {
		String data = username + ":" + tokenExpiryTime + ":" + password + ":" + getKey();
		try {
			MessageDigest digest = MessageDigest.getInstance(algorithm.getDigestAlgorithm());
			return new String(Hex.encode(digest.digest(data.getBytes())));
		} catch (NoSuchAlgorithmException ex) {
			throw new IllegalStateException("No " + algorithm.name() + " algorithm available!");
		}
	}

	/**
	 * cookie 是否超时
	 *
	 * @param tokenExpiryTime 超时时间
	 * @return 是否超时
	 */
	protected boolean isTokenExpired(long tokenExpiryTime) {
		return tokenExpiryTime < System.currentTimeMillis();
	}

	/**
	 * 登录成功处理逻辑
	 *
	 * @param request                  请求都西昂
	 * @param response                 响应对象
	 * @param successfulAuthentication 认证成功封装类
	 */
	@Override
	public void onLoginSuccess(HttpServletRequest request, HttpServletResponse response, Authentication successfulAuthentication) {
		// 获取用户名密码
		String username = retrieveUserName(successfulAuthentication);
		String password = retrievePassword(successfulAuthentication);
		// 用户名密码校验
		if (!StringUtils.hasLength(username)) {
			this.logger.debug("用户名找回失败");
			return;
		}
		if (!StringUtils.hasLength(password)) {
			UserDetails user = getUserDetailsService().loadUserByUsername(username);
			password = user.getPassword();
			if (!StringUtils.hasLength(password)) {
				this.logger.debug("密码找回失败 " + username);
				return;
			}
		}
		// 重置超时时间
		int tokenLifetime = calculateLoginLifetime(request, successfulAuthentication);
		long expiryTime = System.currentTimeMillis();
		expiryTime += 1000L * ((tokenLifetime < 0) ? TWO_WEEKS_S : tokenLifetime);
		// 重新生成签名
		String signatureValue = makeTokenSignature(expiryTime, username, password, this.encodingAlgorithm);
		// 回写到cookie中
		setCookie(new String[]{username, Long.toString(expiryTime), this.encodingAlgorithm.name(), signatureValue},
				tokenLifetime, request, response);
		if (this.logger.isDebugEnabled()) {
			this.logger.debug("Added remember-me cookie for user '" + username + "', expiry: '" + new Date(expiryTime) + "'");
		}
	}

	/**
	 * 设置匹配算法
	 *
	 * @param matchingAlgorithm 匹配算法
	 */
	public void setMatchingAlgorithm(RememberMeTokenAlgorithm matchingAlgorithm) {
		Assert.notNull(matchingAlgorithm, "matchingAlgorithm cannot be null");
		this.matchingAlgorithm = matchingAlgorithm;
	}

	/**
	 * 重置签名存活时间
	 *
	 * @param request        请求对象
	 * @param authentication 加密成功封装对象
	 * @return 签名时长
	 */
	protected int calculateLoginLifetime(HttpServletRequest request, Authentication authentication) {
		return getTokenValiditySeconds();
	}

	/**
	 * 解析用户名
	 *
	 * @param authentication 认证成功封装对象
	 * @return 用户名
	 */
	protected String retrieveUserName(Authentication authentication) {
		if (isInstanceOfUserDetails(authentication)) {
			return ((UserDetails) authentication.getPrincipal()).getUsername();
		}
		return authentication.getPrincipal().toString();
	}

	/**
	 * 解析密码
	 *
	 * @param authentication 认证成功封装对象
	 * @return 密码
	 */
	protected String retrievePassword(Authentication authentication) {
		if (isInstanceOfUserDetails(authentication)) {
			return ((UserDetails) authentication.getPrincipal()).getPassword();
		}
		if (authentication.getCredentials() != null) {
			return authentication.getCredentials().toString();
		}
		return null;
	}

	/**
	 * 判断是否未强制转换字符串
	 * 参见登录成功后强制加密字符串配置
	 *
	 * @param authentication 认证成功封装对象
	 * @return 是否对象
	 */
	private boolean isInstanceOfUserDetails(Authentication authentication) {
		return authentication.getPrincipal() instanceof UserDetails;
	}

	/**
	 * 验签算法
	 *
	 * @param expected 期望签名
	 * @param actual   真实签名
	 * @return 是否验签成功
	 */
	private static boolean equals(String expected, String actual) {
		byte[] expectedBytes = bytesUtf8(expected);
		byte[] actualBytes = bytesUtf8(actual);
		return MessageDigest.isEqual(expectedBytes, actualBytes);
	}

	private static byte[] bytesUtf8(String s) {
		return (s != null) ? Utf8.encode(s) : null;
	}

	/**
	 * 记住我算法枚举类 支持 MD5 SHA-256
	 */
	public enum RememberMeTokenAlgorithm {

		MD5("MD5"), SHA256("SHA-256");

		private final String digestAlgorithm;

		RememberMeTokenAlgorithm(String digestAlgorithm) {
			this.digestAlgorithm = digestAlgorithm;
		}

		public String getDigestAlgorithm() {
			return this.digestAlgorithm;
		}

	}

}
